import json

from fastapi import APIRouter, Body, Query

from server.config import vn
from server.utils.response_util import ResponseUtil
from server.utils.string_util import StringUtil

knController = APIRouter(prefix="/kn", tags=["知识管理"])


@knController.get('/get_kn_data', summary="获取知识数据")
async def get_kn_data(
        kb_type: str = Query(None, description="知识库名称"),
        page_num: int = Query(1, ge=1, alias="pageNum", description="页码"),
        page_size: int = Query(10, alias="pageSize", description="每页数据量")
):
    """
    根据分页参数分割DataFrame。

    :param kb_type: 知识类型
    :param page_num: 当前页码，默认为1
    :param page_size: 每页显示的行数，默认为10
    :return: 分页后的DataFrame和总页数
    """

    df = vn.get_training_data()

    rows = []
    if kb_type:
        df = df[df['training_data_type'] == kb_type]
    total = len(df)

    if total > 0:
        total_pages = (total + page_size - 1) // page_size  # 计算总页数

        if page_num < 1 or page_num > total_pages:
            raise ValueError(f"Page number out of range. Valid pages are between 1 and {total_pages}")

        start = (page_num - 1) * page_size
        end = start + page_size
        page_df = df.iloc[start:end]
        page_df_json_str = page_df.to_json(orient='records', force_ascii=False)
        rows = json.loads(page_df_json_str)

    return ResponseUtil.success(data={
        "rows": rows,
        "total": total
    })


@knController.post("/add_kn_data", summary="添加知识数据")
async def add_kn_data(question: str = Body(None, description="问题"),
                      sql: str = Body(None, description="SQL语句"),
                      ddl: str = Body(None, description="DDL语句"),
                      documentation: str = Body(None, description="文档")):
    try:
        train_id = vn.train(question=question, sql=sql, ddl=ddl, documentation=documentation)

        return ResponseUtil.success(data={"id": train_id})
    except Exception as e:
        print("TRAINING ERROR", e)
        return ResponseUtil.error(msg=str(e))


@knController.post("/remove_kn_data", summary="删除这是数据")
async def remove_kn_data(ids: str = Body(..., embed=True, description="id", example="1,2,3")):
    if StringUtil.is_empty(ids):
        return ResponseUtil.error(msg="No id provided")
    ids_arr = ids.split(",")

    for sub_id in ids_arr:
        vn.remove_training_data(id=sub_id)

    return ResponseUtil.success()
